import numpy as np
import torch
import torch.nn.functional as F

from attack.base_attack import BaseAttack

class Attack(BaseAttack):
    def __init__(self, model, config):
        super().__init__('LLC', model, config, batch_size = 100, targeted = True, llc = True)             #epsilon = 0.99已经基本达到最大的攻击能力，在往上没有更多效果

    def attack_batch(self, images, targets):
        var_images = torch.from_numpy(images).to(self.model.device)
        var_images.requires_grad = True
        var_targets = torch.LongTensor(targets).to(self.model.device)

        self.model.eval()
        output = self.model(var_images)
        loss = F.cross_entropy(output, var_targets)
        loss.backward()
        grad_sign = var_images.grad.sign().cpu().numpy()

        adv_images = images - self.config['epsilon'] * grad_sign
        adv_images = np.clip(adv_images, 0.0, 1.0)

        return adv_images
